# -*- coding: utf-8 -*-
"""
水库调度规程智能问答系统 - Streamlit前端界面
支持流式输出和对话历史
"""

import sys
import os
from pathlib import Path
import streamlit as st
from typing import Generator
import time

# 将项目根目录添加到路径
project_root = Path(__file__).resolve().parent
sys.path.insert(0, str(project_root))

from D_db.tools.llm_query_tester import LLMQueryTester
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

# 页面配置
st.set_page_config(
    page_title="DUT水库调度智能助手",
    page_icon="💧",
    layout="wide",
    initial_sidebar_state="expanded"
)

# 自定义CSS样式
st.markdown("""
<style>
    /* 全局字体：微软雅黑 */
    html, body, .stApp { font-family: "Microsoft YaHei", Arial, sans-serif; }

    .main-header {
        font-family: "SimHei", "Microsoft YaHei", Arial, sans-serif;
        font-size: 24px;
        font-weight: 700;
        color: #000000;
        text-align: center;
        padding: 12px 0;
        margin: 8px 0 16px;
    }
    .chat-message {
        padding: 1rem;
        border-radius: 10px;
        margin-bottom: 1rem;
        box-shadow: 0 2px 4px rgba(0,0,0,0.1);
    }
    .user-message {
        background-color: #E3F2FD;
        border-left: 4px solid #1E88E5;
        font-size: 14px; /* 问题文本：14号常规 */
        font-weight: 400;
        line-height: 1.7;
    }
    .assistant-message {
        background-color: #F5F5F5;
        border-left: 4px solid #4CAF50;
        font-size: 16px; /* 答案文本：16号加粗 */
        font-weight: 700;
        line-height: 1.7;
    }
    .stTextInput > div > div > input {
        font-size: 1.1rem;
    }
    .info-box {
        background-color: #FFF3E0;
        padding: 1rem;
        border-radius: 8px;
        border-left: 4px solid #FF9800;
        margin: 1rem 0;
    }

    /* 分栏标题样式 */
    .section-title {
        font-size: 16px;
        font-weight: 700;
        margin: 0 0 8px 0;
    }

    /* 两列布局的栅格间距在小屏下收敛 */
    @media (max-width: 992px) {
        .chat-message {
            margin-bottom: 0.75rem;
        }
    }
    
    /* 聊天界面样式优化 */
    .stChat {
        height: 60vh;
        overflow-y: auto;
        padding: 20px;
        background-color: #fafafa;
        border-radius: 10px;
        margin-bottom: 20px;
    }
    
    /* 用户消息气泡样式 */
    .user-bubble {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        color: white;
        padding: 12px 16px;
        border-radius: 18px 18px 4px 18px;
        margin: 8px 0;
        max-width: 70%;
        margin-left: auto;
        box-shadow: 0 2px 8px rgba(102, 126, 234, 0.3);
        font-size: 14px;
        line-height: 1.5;
    }
    
    /* 助手消息气泡样式 */
    .assistant-bubble {
        background: white;
        color: #333;
        padding: 12px 16px;
        border-radius: 18px 18px 18px 4px;
        margin: 8px 0;
        max-width: 70%;
        margin-right: auto;
        box-shadow: 0 2px 8px rgba(0, 0, 0, 0.1);
        border: 1px solid #e0e0e0;
        font-size: 15px;
        line-height: 1.6;
    }
    
    /* 输入框样式优化 */
    .stChatInput {
        background: white;
        border-radius: 25px;
        border: 2px solid #e0e0e0;
        padding: 12px 20px;
        font-size: 16px;
        transition: all 0.3s ease;
    }
    
    .stChatInput:focus {
        border-color: #667eea;
        box-shadow: 0 0 0 3px rgba(102, 126, 234, 0.1);
        outline: none;
    }
    
    /* 滚动条样式 */
    ::-webkit-scrollbar {
        width: 6px;
    }
    
    ::-webkit-scrollbar-track {
        background: #f1f1f1;
        border-radius: 3px;
    }
    
    ::-webkit-scrollbar-thumb {
        background: #c1c1c1;
        border-radius: 3px;
    }
    
    ::-webkit-scrollbar-thumb:hover {
        background: #a8a8a8;
    }
</style>
""", unsafe_allow_html=True)

def initialize_session_state():
    """初始化 session state"""
    if "messages" not in st.session_state:
        st.session_state.messages = []
    
    if "tester" not in st.session_state:
        try:
            st.session_state.tester = LLMQueryTester()
            st.session_state.tester_initialized = True
        except Exception as e:
            st.session_state.tester = None
            st.session_state.tester_initialized = False
            st.session_state.init_error = str(e)

def stream_response(response_text: str) -> Generator[str, None, None]:
    """模拟流式输出效果"""
    words = response_text.split()
    for i, word in enumerate(words):
        if i > 0:
            yield " "
        yield word
        time.sleep(0.05)  # 控制输出速度

def format_qa_results(qa_results: list) -> str:
    """格式化查询结果"""
    if not qa_results:
        return ""
    
    result_text = "\n\n---\n\n### 📚 参考的知识库内容\n\n"
    
    for i, item in enumerate(qa_results[:5], 1):  # 只显示前5个
        result_text += f"**{i}. {item.get('question', '')}**\n\n"
        result_text += f"*答案：{item.get('answer', '')}*\n\n"
        
        if item.get('category_id'):
            result_text += f"🏷️ 水库ID: `{item['category_id']}`  "
        if item.get('match_count'):
            result_text += f"| 🎯 匹配度: `{item['match_count']}`\n\n"
        
        result_text += "---\n\n"
    
    return result_text

def display_message(role: str, content: str, qa_results: list = None):
    """显示消息"""
    if role == "user":
        st.markdown(f"""
        <div class="chat-message user-message">
            <strong>🙋 您的问题：</strong><br>
            {content}
        </div>
        """, unsafe_allow_html=True)
    else:
        st.markdown(f"""
        <div class="chat-message assistant-message">
            <strong>🤖 智能助手：</strong><br>
        </div>
        """, unsafe_allow_html=True)
        
        # 流式显示回答
        if content:
            st.markdown(content)
        
        # 显示参考资料
        if qa_results:
            with st.expander("🔍 查看参考的知识库内容", expanded=False):
                formatted_results = format_qa_results(qa_results)
                st.markdown(formatted_results)

def render_chat_message(message, index):
    """渲染单个聊天消息，根据角色决定左右对齐"""
    if message["role"] == "user":
        # 用户消息：右侧显示
        cols = st.columns([1, 6])
        with cols[1]:
            st.markdown(f"""
            <div style="text-align: right; margin: 8px 0;">
                <div class="user-bubble">
                    {message['content']}
                </div>
            </div>
            """, unsafe_allow_html=True)
    else:
        # 助手消息：左侧显示
        cols = st.columns([6, 1])
        with cols[0]:
            st.markdown(f"""
            <div style="text-align: left; margin: 8px 0;">
                <div class="assistant-bubble">
                    <div style="margin-bottom: 8px;"><strong>🤖 智能助手</strong></div>
                    {message['content']}
                </div>
            </div>
            """, unsafe_allow_html=True)
            
            # 显示参考资料
            if message.get("qa_results"):
                with st.expander("🔍 查看参考的知识库内容", expanded=False):
                    formatted_results = format_qa_results(message["qa_results"])
                    st.markdown(formatted_results)

def render_chat_conversation():
    """渲染聊天对话，按时间顺序显示消息"""
    if not st.session_state.messages:
        st.markdown("""
        <div style="text-align: center; color: #666; margin: 100px 0;">
            <div style="font-size: 48px; margin-bottom: 20px;">💧</div>
            <h2 style="color: #333; margin-bottom: 10px;">DUT水库调度智能助手</h2>
            <p style="font-size: 16px; color: #666;">请在下方输入您的问题，我将为您提供专业的解答</p>
            <div style="margin-top: 30px;">
                <span style="background: #f0f0f0; padding: 8px 16px; border-radius: 20px; font-size: 14px;">
                    💡 支持自然语言提问
                </span>
            </div>
        </div>
        """, unsafe_allow_html=True)
        return
    
    # 创建聊天容器，带滚动效果
    chat_container = st.container()
    
    with chat_container:
        # 添加聊天头部
        st.markdown("""
        <div style="text-align: center; margin-bottom: 20px; padding: 10px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); border-radius: 10px; color: white;">
            <strong>💬 智能对话</strong>
        </div>
        """, unsafe_allow_html=True)
        
        # 按时间顺序显示所有消息
        for i, message in enumerate(st.session_state.messages):
            render_chat_message(message, i)
            
        # 添加一些间距
        st.markdown("<div style='height: 30px;'></div>", unsafe_allow_html=True)

def main():
    """主函数"""
    # 初始化
    initialize_session_state()
    
    # 标题（加粗24号黑体，居中）
    st.markdown('<h1 class="main-header">DUT水库调度智能助手</h1>', unsafe_allow_html=True)
    
    # 侧边栏
    with st.sidebar:
        st.header("📜 对话历史")
        
        # 检查系统状态
        if st.session_state.tester_initialized:
            # 对话历史记录
            if st.session_state.messages:
                # 显示对话历史
                for i, message in enumerate(st.session_state.messages):
                    if message["role"] == "user":
                        st.markdown(f"**👤 问题 {i//2 + 1}:**")
                        st.markdown(f"{message['content']}")
                    else:
                        st.markdown(f"**🤖 回答 {i//2 + 1}:**")
                        # 只显示回答的前100个字符，避免太长
                        answer_preview = message['content'][:100] + "..." if len(message['content']) > 100 else message['content']
                        st.markdown(f"{answer_preview}")
                        if len(message['content']) > 100:
                            if st.button("显示完整回答", key=f"show_full_{i}"):
                                st.session_state[f"show_full_{i}"] = True
                        if st.session_state.get(f"show_full_{i}", False):
                            st.markdown(f"{message['content']}")
                    st.markdown("---")
            else:
                st.markdown("*暂无对话记录*")
            
            # 清空对话按钮
            if st.button("🗑️ 清空对话历史", use_container_width=True):
                st.session_state.messages = []
                st.rerun()
        else:
            st.error("❌ 系统未就绪")
            st.error(f"错误信息：{st.session_state.get('init_error', '未知错误')}")
            
            if st.button("🔄 重新初始化", use_container_width=True):
                if "tester" in st.session_state:
                    del st.session_state.tester
                if "tester_initialized" in st.session_state:
                    del st.session_state.tester_initialized
                st.rerun()
    
    # 主界面
    if not st.session_state.tester_initialized:
        st.warning("⚠️ 系统未就绪，请查看左侧边栏的错误信息")
        return
    
    # 显示对话（DeepSeek风格布局）
    render_chat_conversation()
    
    # 处理选中的示例问题
    if hasattr(st.session_state, 'selected_question'):
        user_question = st.session_state.selected_question
        del st.session_state.selected_question
    else:
        user_question = None
    
    # 输入框
    with st.container():
        col1, col2 = st.columns([6, 1])
        
        with col1:
            prompt = st.chat_input(
                "请输入您的问题...",
                key="user_input"
            )
        
        with col2:
            st.markdown("<br>", unsafe_allow_html=True)
    
    # 处理用户输入
    if prompt or user_question:
        question = user_question if user_question else prompt
        
        # 添加用户消息（右侧展示）
        st.session_state.messages.append({"role": "user", "content": question})
        
        # 显示正在处理的提示
        with st.spinner("🔍 正在查询知识库并生成回答..."):
            try:
                # 获取查询结果
                limit = st.session_state.get('query_limit', 20)
                result = st.session_state.tester.test_query(question, limit=limit)
                
                answer = result.get("answer", "抱歉，我无法回答这个问题。")
                qa_results = result.get("qa_results", [])
                
                # 添加助手消息（左侧展示）
                st.session_state.messages.append({
                    "role": "assistant",
                    "content": answer,
                    "qa_results": qa_results
                })
                
            except Exception as e:
                error_msg = f"❌ 查询出错：{str(e)}"
                st.error(error_msg)
                st.session_state.messages.append({
                    "role": "assistant",
                    "content": error_msg
                })
        
        # 重新运行以更新界面
        st.rerun()
    
    # 添加一些示例问题按钮（可选）
    if not st.session_state.messages:
        st.markdown("""
        <div style="text-align: center; margin-top: 20px;">
            <p style="color: #666; font-size: 14px; margin-bottom: 10px;">试试这些问题：</p>
        </div>
        """, unsafe_allow_html=True)
        
        # 示例问题按钮
        sample_questions = [
            "大伙房水库的调度规则是什么？",
            "柴河水库的汛期是什么时候？",
            "汤河水库的正常蓄水位是多少？"
        ]
        
        cols = st.columns(len(sample_questions))
        for i, question in enumerate(sample_questions):
            with cols[i]:
                if st.button(f"{question}", key=f"sample_{i}", use_container_width=True):
                    st.session_state.selected_question = question
                    st.rerun()

if __name__ == "__main__":
    main()

